#!/usr/bin/python
#Created by David Langelaan 2008
#Filename: SparkyResonanceList.py
'''A sript which takes an amino acid sequence, and then calculates the complete
sparky resonance list which you will need to assign. This list is saved as a txt file
which can be imported into sparky'''

#Insert the sequence of your peptide here...no spaces
sequence = 'CNESGFSWGYPHPFLI'

#This is the output file directory where you want your file to go, 'file.txt'
#will just put a file in the current working directory, unless you specify a directory structure
Outputfilename = '/Users/Dave/Desktop/out.txt'


# A dictionary of all of the residues supported, feel free to add...Cystein defaults to reduced form, proline to trans..
# Random shifts from Wishart et al, Journal of Biomolecular NMR, 5(1995) 67-81. I guessed on arg HH, Tyr HH , His HE2 to be 0.00 since they are not in Wishart et al. 
Residues = {'A': [['HN', 8.24], ['HA', 4.32], ['HB', 1.39], ['CA', 52.5], ['CB', 19.1]],
            'C': [['HN', 8.32], ['HA', 4.55], ['HB', 2.93], ['CA', 58.2], ['CB', 28.0]],
            'D': [['HN', 8.34], ['HA', 4.64], ['HBa', 2.72], ['HBb', 2.65], ['CA', 54.2], ['CB', 41.1]],
            'E': [['HN', 8.42], ['HA', 4.35], ['HBa', 2.06], ['HBb', 1.96], ['HG', 2.31], ['CA', 56.6], ['CB', 29.9], ['CG', 35.6]],
            'F': [['HN', 8.30], ['HA', 4.62], ['HBa', 3.14], ['HBb', 3.04], ['HD', 7.28], ['HE', 7.38], ['HZ', 7.32], ['CA', 57.7], ['CB', 39.6], ['CD', 131.9], ['CE', 131.5], ['CZ', 129.9]],
            'G': [['HN', 8.33], ['HA', 3.96], ['CA', 45.1]],
            'H': [['HN', 8.42], ['HA', 4.73], ['HBa', 3.29], ['HBb', 3.16], ['HD2', 7.29], ['HE1',8.58], ['HE2', 0.00], ['CA', 55.0], ['CB', 29.0], ['CD2', 120.1], ['CE1', 136.2]],
            'I': [['HN', 8.00], ['HA', 4.17], ['HB', 1.87], ['HG11', 1.45], ['HG12', 1.16], ['HG2', 0.91], ['HD', 0.86], ['CA', 61.1], ['CB', 38.8], ['CG1', 27.2], ['CG2', 17.4], ['CD', 12.9]], 
            'K': [['HN', 8.29], ['HA', 4.32], ['HBa', 1.84], ['HBb', 1.75], ['HG', 1.44], ['HD', 1.68], ['HE', 2.99], ['HZ', 7.81], ['CA', 56.2], ['CB', 33.1], ['CG', 24.7], ['CD', 29.0], ['CE', 41.9]],
            'L': [['HN', 8.16], ['HA', 4.34], ['HB', 1.62], ['HG', 1.59], ['HD1', 0.92], ['HD2', 0.87], ['CA', 55.1], ['CB', 42.4], ['CG', 26.9], ['CD1', 24.9], ['CD2', 23.3]],
            'M': [['HN', 8.28], ['HA', 4.48], ['HBa', 2.11], ['HBb', 2.01], ['HGa',2.60], ['HGb', 2.54], ['HE', 2.10], ['CA', 55.4], ['CB', 32.9], ['CG', 32.0], ['CE', 16.9]],
            'N': [['HN', 8.40], ['HA', 4.74], ['HBa', 2.83], ['HBb', 2.75], ['HD2a', 7.59], ['HD2b', 6.91], ['CA', 53.1], ['CB', 38.9]],
            'P': [['HA', 4.42], ['HBa', 2.29], ['HBb', 1.94], ['HG', 2.02], ['HD', 3.63], ['CA', 63.3], ['CB', 32.1], ['CG', 27.2], ['CD', 49.8]],
            'Q': [['HN', 8.32], ['HA', 4.34], ['HBa', 2.12], ['HBb', 1.99], ['HG', 2.36], ['HE2a', 7.52], ['HE2b', 6.85], ['CA', 55.7], ['CB', 29.4], ['CG', 33.7]],
            'R': [['HN', 8.23], ['HA', 4.34], ['HBa', 1.86], ['HBb', 1.76], ['HG', 1.63], ['HD', 3.20], ['HE', 8.07], ['HH', 0.00], ['CA', 56.0], ['CB', 30.9], ['CG', 27.1], ['CD', 43.3]],
            'S': [['HN', 8.31], ['HA', 4.47], ['HBa', 3.89], ['HBb', 3.87], ['CA', 58.3], ['CB', 63.8]],
            'T': [['HN', 8.15], ['HA', 4.35], ['HB', 4.24], ['HG', 1.21], ['CA', 61.8], ['CB', 69.8], ['CG', 21.5]],
            'V': [['HN', 8.03], ['HA', 4.12], ['HB', 2.08], ['HG1', 0.94], ['HG2', 0.93], ['CA', 62.2], ['CB', 32.9], ['CG1', 21.1], ['CG2', 20.3]],
            'W': [['HN', 8.25], ['HA', 4.66], ['HBa', 3.29], ['HBb', 3.27],['HD1', 7.27], ['HE3', 7.65], ['HZ3', 7.18], ['HH2', 7.25], ['HZ2', 7.50], ['CA', 57.5], ['CB', 29.6], ['CD1', 127.4], ['CE3', 122.2], ['CZ3', 124.8], ['CH2', 121.0], ['CZ2', 114.7]],
            'Y': [['HN', 8.12], ['HA', 4.55], ['HBa', 3.03], ['HBb', 2.98], ['HD', 7.14], ['HE', 6.84],['HH', 0.00], ['CA', 57.9], ['CB', 38.8], ['CD', 133.3], ['CE', 118.2]]}


###The business end of the program....
f = open(Outputfilename, 'w')
print >> f, '%-5s %-5s %-5s %-5s' %('Group', 'Atom', 'Nuc', 'Shift')
for index, residue in enumerate(sequence):
    for aminoacid in Residues[residue]:
        if aminoacid[0][0]== 'H':
            Nuc = '1H'
        if aminoacid[0][0]== 'C':
            Nuc = '13C'
        print >> f, '%-5s %-5s %-5s %.2f' %(residue+str(index+1), aminoacid[0], Nuc, aminoacid[1])
f.close()

print 'Resonancelist generated successfully, now get to work!'


#another part of the program, uses the same dictionary and sequence as earlir.Creates tables which you can fill in as you analyze a structure.
for index, residue in enumerate(sequence):
    width = 1
    Hwidth = 1
    print 'TOCSY Connections within %s%s' %(residue, index+1)
    print '%s%-4s' %(residue, index+1),
    for atom in Residues[str(residue)]:
        print '%-4s' %atom[0],
        if atom[0][0] == 'H':
            Hwidth +=1
    print ' '
    for atom in Residues[str(residue)]:
            if atom[0][0] =='H':
                print '%-4s' %atom[0],
                for number in range(0,Hwidth-1):
                    print '%-4s' %'|',
                print '%-4s' %'|'
    #makes chart for 1+1 connectivities
    try:
        print 'NOESY i+1 Connections to %s%s' %(sequence[index+1], index+2)
        for atom in Residues[str(sequence[index+1])]:
            if atom[0][0] =='H':
                print '%-4s' %atom[0],
                for number in range(0,Hwidth-1):
                    print '%-4s' %'|',
                print '%-4s' %'|'
    except IndexError:
        print 'No i+1 Connections'
    #makes chart for i+2 connectivities
    try:
        print 'NOESY i+2 Connections to %s%s' %(sequence[index+2], index+3)
        for atom in Residues[str(sequence[index+2])]:
            if atom[0][0] == 'H':
                print '%-4s' %atom[0],
                for number in range(0,Hwidth-1):
                    print '%-4s' %'|',
                print '%-4s' %'|'
    except IndexError:
        print 'No i+2 Connections'
    print '\n\n\n'

##Another part of the program, which makes a shortened worksheet for resonance assignment from HSQC, TOCSY, NOESY
##Making table for TOCSY assignment

#a function which prints lists in the form of tables
def printtable(data):
  counter =0
  while True:
    for line in data:
      i = 0+counter*10
      if i >= len(sequence): break
      while i < counter*10+10:   
        print '%10s' %line[i],
        i+=1
        if i == len(sequence):
          break
      print ''
    print '\n\n'
    counter +=1
    if counter*10 > len(sequence):
      break
#a function which creates a list which is in a good format for printtable
def maketable(resonance):
  nResidues = {}
  for residue in sequence:
    nResidues[residue] =[]
    for atom in Residues[residue]:
      if atom[0][0] == resonance:
        nResidues[residue].append(atom)
  maxats = 0
  lines = [[]]
  for index, residue in enumerate(sequence):
    ats = 0
    lines[0].append( residue +  str(index+1) + '   |')
    for atom in nResidues[residue]:
      if resonance in atom[0]:
        ats +=1
    if ats > maxats:
      maxats = ats

  for i in range(0,maxats):
    lines.append([])
    for residue in sequence:
      try:
        if resonance in nResidues[residue][i][0][0]:
          a = '%10s' %(str((nResidues[residue][i][0] +':'+ str(nResidues[residue][i][1])+ '|')))
          lines[i+1].append(a)
        if not resonance in nResidues[residue][i][0][0]:
          lines[i+1].append('         |')
      except:
        lines[i+1].append('         |')
  #print '\n\n\n'
  #print nResidues
  return lines


print 'Hydrogen Resonances from TOCSY'
lines = maketable('H')
printtable(lines)

print 'Carbon Resonances from HSQC'
lines = maketable('C')
printtable(lines)

print 'NOESY Assignment Summary'
data = [['Residue'],['Intra'],['i+1'],['i+2'],['i+3']]

#generating the table for the NOESY part(it is a different type of table)
for index, residue in enumerate(sequence): 
  data[0].append('|'+residue +str(index+1))
i=0
while i < len(sequence):
  for row in data[1:]:
    row.append('|')
  i+=1
#printing the table for the NOESY part
i =1
while i < len(sequence):
  for row in data:
    print '%-7s' %(row[0]),
    for item in row[i:i+20]:
      print '%-4s' %item,
    print ''
  print '\n\n'
  i+=20


#for residue in sequence:
#  print residue
#  for atom in Residues[residue]:
#    print atom[0], atom[1],
#  print '\n'
